"""
@File    : mysql_config.py
@Author  : timo.tan
@Time    : 2025/07/23 17:17 
"""

from sqlalchemy import create_engine, text
from sqlalchemy.pool import QueuePool
from sqlalchemy.orm import sessionmaker
import logging
from contextlib import contextmanager


class MySQLClient:
    """MySQL数据库客户端封装类"""

    def __init__(self, host, port, user, password, database, pool_size=10, max_overflow=20):
        """
        初始化MySQL连接
        
        Args:
            host: 数据库主机地址
            port: 数据库端口
            user: 用户名
            password: 密码
            database: 数据库名
            pool_size: 连接池大小
            max_overflow: 最大溢出连接数
        """
        self.connection_url = f"mysql+pymysql://{user}:{password}@{host}:{port}/{database}"
        self.engine = create_engine(
            self.connection_url,
            poolclass=QueuePool,
            pool_size=pool_size,
            max_overflow=max_overflow,
            pool_pre_ping=True,
            pool_recycle=3600
        )
        self.Session = sessionmaker(bind=self.engine)
        self.logger = logging.getLogger(__name__)

    @contextmanager
    def get_session(self):
        """获取数据库会话的上下文管理器"""
        session = self.Session()
        try:
            yield session
            session.commit()
        except Exception as e:
            session.rollback()
            self.logger.error(f"数据库操作失败: {str(e)}")
            raise
        finally:
            session.close()

    def execute_query(self, query, params=None):
        """
        执行查询SQL
        
        Args:
            query: SQL查询语句
            params: 参数字典
            
        Returns:
            查询结果值
        """
        with self.get_session() as session:
            cursor_result = session.execute(text(query), params or {})
            row = cursor_result.fetchall()
            if row is None:
                return None

            # 如果只有一列，直接返回值
            if len(row) == 1:
                return row[0]
            return row

    def execute_update(self, query, params=None):
        """
        执行更新SQL
        
        Args:
            query: SQL更新语句
            params: 参数字典
            
        Returns:
            影响的行数
        """
        with self.get_session() as session:
            result = session.execute(text(query), params or {})
            # session.commit() 在上下文管理器中自动调用
            return result.rowcount


db = MySQLClient(host='10.141.194.72', port=23306, user='viu_root', password='myviu_4359', database='d_cgboss_info')

if __name__ == '__main__':
    project_id = db.execute_query("select instance_id, project_id from iaas_vm ")
    print(project_id)


